package utils

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"encoding/base64"

	log "github.com/sirupsen/logrus"
)

func pkCS5Unpadding(origData []byte) []byte {
	length := len(origData)
	unpadding := int(origData[length-1])
	return origData[:(length - unpadding)]
}

func Encrypt(origin_data string, key string) (base64_result string, err error) {
	iv := []byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6}
	// fmt.Println(string(iv))
	// iv := []byte("0000000000000000")
	var block cipher.Block
	if block, err = aes.NewCipher([]byte(key)); err != nil {
		log.Println(err)
		return
	}
	encrypt := cipher.NewCBCEncrypter(block, iv)
	var source []byte = PKCS5Padding([]byte(origin_data), 16)
	var dst []byte = make([]byte, len(source))
	encrypt.CryptBlocks(dst, source)
	base64_result = base64.StdEncoding.EncodeToString(dst)
	return
}

func PKCS5Padding(ciphertext []byte, blockSize int) []byte {
	padding := blockSize - len(ciphertext)%blockSize
	padtext := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(ciphertext, padtext...)
}

func DecodeCBC(encrypt_data string, key string) (origin_data string, err error) {

	defer func() {
		if err := recover(); err != nil {
			log.Errorf("decode error: %v", err)
		}
	}()
	var source []byte
	if encrypt_data[len(encrypt_data)-1] == '\n' {
		encrypt_data = encrypt_data[:len(encrypt_data)-1]
	}
	if source, err = base64.StdEncoding.DecodeString(encrypt_data); err != nil {
		log.Errorf("decode base64 error: %v", err)
		return
	}

	iv := []byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6}
	var block cipher.Block
	if block, err = aes.NewCipher([]byte(key)); err != nil {
		log.Errorf("new cipher error: %v", err)
		return
	}
	encrypt := cipher.NewCBCDecrypter(block, iv)

	var dst []byte = make([]byte, len(source))
	encrypt.CryptBlocks(dst, source)
	origin_data = string(pkCS5Unpadding(dst))
	return
}
